# -*- coding: utf-8 -*-
# time: 2025/3/27 16:56
# file: chromadb_ollama_ARG.py
# author: hanson
"""
https://www.datacamp.com/tutorial/deepseek-r1-rag
https://article.juejin.cn/post/7470807715898212406
"""
from langchain_ollama import OllamaEmbeddings, ChatOllama
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_community.document_loaders import TextLoader
import os
import shutil

# 1. 加载文档（示例使用单个文本文件）
VECTOR_DIR = "data/vector_store"
# 3. 初始化模型 deepseek-r1:7b qwen2.5:7b
embeddings = OllamaEmbeddings(model="qwen2.5:7b")
llm = ChatOllama(model="qwen2.5:7b", temperature=0.3)

# 4. 创建向量数据库（自动持久化）


vector_db = Chroma(
    #documents=split_docs,
    persist_directory=VECTOR_DIR,
    embedding_function=embeddings
)
print("向量数据库已创建")

# 5. 构建检索链（优化中文模板）

prompt_template = """
您是一个设计用于査询文档来回答问题的代理。您可以使用文档检索工具。
并基于检索内容来回答问题您可能不查询文档就知道答案，但是您仍然应该查询文档来获得答案。
如果您从文档中找不到任何信息用于回答问题，则只需返回“抱歉，这个问题我还不知道。”作为答案。
如果有人提问等关于您的名字的问题，您就回答：“我是超级牛逼哄哄的小天才助手”作为答案。
上下文：{context}
问题：{question}
"""

retriever = vector_db.as_retriever(
    search_type="mmr",  # 最大边际相关性搜索（平衡相关性与多样性）
    search_kwargs={
        "k": 2,  # 初始检索文档数量 ，精度高点可以调小
        "fetch_k": 20,  # 底层搜索数量（越大精度越高）
        "lambda_mult": 0.5,  # 多样性控制参数（0-1，越大越多样）
        "score_threshold": 0.3,  # 相关性阈值
    },
)

chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | ChatPromptTemplate.from_template(prompt_template)
    | llm
)

# 6. 执行查询
response = chain.invoke("我是百度地图开放平台认证开发者，是不是代表我在其他平台均可以此身份作为身份说明？")
print(f"\n回答：\n{response.content}")

